/*
  This file is part of bnc.

  bnc is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  bnc is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with bnc.  If not, see <http://www.gnu.org/licenses/>.
*/
#include "./compressor.h"

void compressorConstructor (compressor* aCompressor)
{
  /*
    Prepares stack and sensitive variables.
  */
  aCompressor->pStack=(stack*)malloc(sizeof(stack));
  stackConstructor(aCompressor->pStack);

  aCompressor->pPhase=0;
  aCompressor->pTable=0;
  aCompressor->pTree=0;
  aCompressor->pNode=0;
  aCompressor->pLength=0;
}

void compressorReadMatrix (compressor* aCompressor, unsigned char* aBufferI, unsigned int aLength)
{
  /*
    Reads file, counts weights of characters.
  */
  unsigned int i;

  if (aCompressor->pLength==0)  //If this is the first read, prepare leaves on the stack.
  {
    for (i=255; i+1>0; i--)
    {
      stackPush(aCompressor->pStack, (node*)malloc(sizeof(node)));
      nodeConstructor(stackTop(aCompressor->pStack, 0), i);
    }
  }

  aCompressor->pLength+=aLength;

  for (i=0; i<aLength; i++)  //Read the buffer and count the appearances of the characters.
  {
    nodeSetWeight(stackTop(aCompressor->pStack, aBufferI[i]), nodeGetWeight(stackTop(aCompressor->pStack, aBufferI[i]))+1);
  }
}

unsigned int compressorBuildHeader (compressor* aCompressor, unsigned char* aBufferO)
{
  /*
    Builds binary tree, builds translation table.
  */
  unsigned int i=255;

  memcpy(aBufferO, &aCompressor->pLength, 4);  //Header of the compressed file starts with the length of the original file.

  stackSort(aCompressor->pStack, 256);  //Sort the stack so that the leaves with weight zero are at the bottom.

  while (nodeGetWeight(stackTop(aCompressor->pStack, i))==0) i--;  //Skip the leaves that have weight equal to zero.

  while (i>0)  //Connect two least weighted leaves with one node, put that node on the stack, and sort the stack.
  {
    stackPush(aCompressor->pStack, nodeJoin(stackTop(aCompressor->pStack, i), stackTop(aCompressor->pStack, i-1)));
    i--;
    stackSort(aCompressor->pStack, i);
  }

  aCompressor->pTable=(char**)malloc(256*sizeof(char*));  //Prepare table for translation values.

  for (i=0; i<256; i++)  //Get translation for every byte.
  {
    aCompressor->pPhase=-1;
    aCompressor->pTable[i]=(char*)malloc(34*sizeof(char));  //In the worst case the translation string can be 34 bytes long.

    nodeSearch(stackTop(aCompressor->pStack, 0), (unsigned char)i, &aCompressor->pPhase, aCompressor->pTable[i]+1);
    aCompressor->pTable[i][0]=aCompressor->pPhase+1;

    aCompressor->pTable[i][aCompressor->pTable[i][0]/8+1]&=~(255>>(aCompressor->pTable[i][0]&7));  //Cut the mess thats right after the translation.
    aCompressor->pPhase=-1;
  }

  aCompressor->pPhase=0;

  compressorBuildBinaryTree(aCompressor, stackTop(aCompressor->pStack, 0), aBufferO+4);  //Build the binary tree.

  aCompressor->pLastByte=aBufferO[aCompressor->pPhase/8+4];  //Save last byte to ensure continuity of compression as it is bitwise and the compressed data bits may share byte with the end of the tree.

  return aCompressor->pPhase/8+4;
}

void compressorBuildBinaryTree (compressor* aCompressor, node* aNode, unsigned char* aBufferO)
{
  /*
    Encodes binary tree.
    Goes throughout the tree recursively.
  */
  if (nodeIsLeaf(aNode))  //If its leaf write one and store the value as well.
  {
    aBufferO[aCompressor->pPhase/8]|=128>>(aCompressor->pPhase&7);

    aCompressor->pPhase++;

    aBufferO[aCompressor->pPhase/8]|=nodeGetValue(aNode)>>(aCompressor->pPhase&7);
    aBufferO[aCompressor->pPhase/8+1]|=nodeGetValue(aNode)<<(8-(aCompressor->pPhase&7));

    aCompressor->pPhase+=8;
  } else {  //If its node write zero.
    node* lChild;

    aCompressor->pPhase++;

    lChild=nodeGetChildL(aNode);  //Recursively follow the left edge.
    compressorBuildBinaryTree(aCompressor, lChild, aBufferO);

    lChild=nodeGetChildR(aNode);  //Recursively follow the right edge.
    compressorBuildBinaryTree(aCompressor, lChild, aBufferO);
  }
}

unsigned int compressorCompress (compressor* aCompressor, unsigned char* aBufferI, unsigned int aLength, unsigned char* aBufferO)
{
  /*
    Reads and translates input buffer into output buffer.
  */
  unsigned int i;

  aCompressor->pPhase&=7;  //At the beginning we need to know only the offset inside the byte (modulo 8 is equal to bitwise and 7).
  aBufferO[0]=aCompressor->pLastByte;  //Recover last byte.

  for (i=0; i<aLength; i++)  //For every character in the buffer transfer its transletion into output buffer.
  {
    unsigned int j;

    for (j=0; j<(aCompressor->pTable[(unsigned char)aBufferI[i]][0])/8+1; j++)  //Transfer of the translation (byte by byte).
    {
      aBufferO[aCompressor->pPhase/8+j]|=((unsigned char)aCompressor->pTable[(unsigned char)aBufferI[i]][j+1]>>(aCompressor->pPhase&7));
      aBufferO[aCompressor->pPhase/8+j+1]|=((unsigned char)aCompressor->pTable[(unsigned char)aBufferI[i]][j+1]<<(8-(aCompressor->pPhase&7)));
    }

    aCompressor->pPhase+=(unsigned char)aCompressor->pTable[(unsigned char)aBufferI[i]][0];
  }

  if ((aCompressor->pPhase&7)==0) aCompressor->pLastByte=0;  //If the last byte is filled entirely we dont store any value.
  else aCompressor->pLastByte=aBufferO[aCompressor->pPhase/8];  //Otherwise we have to ensure continuity.

  return aCompressor->pPhase/8;  //Return how many bytes the translation took.
}

unsigned char compressorGetLastByte(compressor* aCompressor)
{
  /*
    Gets last byte from last compression to provide continuity.
  */
  aCompressor->pPhase&=7;
  return aCompressor->pLastByte;
}

unsigned char compressorReadBit (compressor* aCompressor, char* aBufferI)
{
  /*
    Gets the value of given bit.
  */
  unsigned char lRet=(aBufferI[aCompressor->pPhase/8]&(128>>(aCompressor->pPhase&7)))>>(7-aCompressor->pPhase&7);
  aCompressor->pPhase++;

  return lRet;
}

char compressorReadByte (compressor* aCompressor, char* aBufferI)
{
  /*
    Recovers byte value at given position.
  */
  char lRet;

  lRet=((unsigned char)aBufferI[aCompressor->pPhase/8])<<(aCompressor->pPhase&7);
  lRet|=((unsigned char)aBufferI[aCompressor->pPhase/8+1])>>(8-(aCompressor->pPhase&7));

  aCompressor->pPhase+=8;

  return lRet;
}

void compressorRecoverHeader (compressor* aCompressor, char* aBufferI)
{
  /*
    Recovers compression header.
    Creates compression tree.
  */
  unsigned int i;

  memcpy(&aCompressor->pLength, aBufferI, 4);  //Recover the original file length.

  aCompressor->pPhase+=32;  //Shift the reading head over the start of the binary tree.

  aCompressor->pTree=(node**)malloc(512*sizeof(node*));
  memset(aCompressor->pTree, 0, 512*sizeof(node*));  //Initialize the array for tree nodes and leaves.

  if (compressorReadBit(aCompressor, aBufferI))  //If the first node is leaf as well we make only one node.
  {
    aCompressor->pTree[0]=(node*)malloc(sizeof(node));
    nodeConstructor(aCompressor->pTree[0], compressorReadByte(aCompressor, aBufferI));
  } else {  //Otherwise we build whole tree.
    aCompressor->pTree[0]=(node*)malloc(sizeof(node));
    nodeConstructor(aCompressor->pTree[0], 0);

    stackPush(aCompressor->pStack, aCompressor->pTree[0]);

    i=1;

    while (stackSize(aCompressor->pStack)!=0)  //If there is still a node without exactly two children then continue reading.
      if (compressorReadBit(aCompressor, aBufferI))  //If there is a leaf then make it child of the node on the top of the stack.
      {
        aCompressor->pTree[i]=(node*)malloc(sizeof(node));
        nodeConstructor(aCompressor->pTree[i], compressorReadByte(aCompressor, aBufferI));
        nodeInjectChild(stackTop(aCompressor->pStack, 0), aCompressor->pTree[i]);
        if (nodeCountChildren(stackTop(aCompressor->pStack, 0))==2) stackPop(aCompressor->pStack);  //If the node on top has got two children, erase it from the stack.
        i++;
      } else {  //If there is a node then put it on top of the stack.
        aCompressor->pTree[i]=(node*)malloc(sizeof(node));
        nodeConstructor(aCompressor->pTree[i], 0);
        nodeInjectChild(stackTop(aCompressor->pStack, 0), aCompressor->pTree[i]);
        if (nodeCountChildren(stackTop(aCompressor->pStack, 0))==2) stackPop(aCompressor->pStack);  //If the node on top has got two children, erase it from the stack.
        stackPush(aCompressor->pStack, aCompressor->pTree[i]);
        i++;
      }
  }
}

unsigned int compressorDecompress (compressor* aCompressor, char* aBufferI, unsigned int aLength, char* aBufferO)
{
  /*
    Decompresses data.
  */
  unsigned int i, lLength=0;

  if (aCompressor->pNode==0) aCompressor->pNode=aCompressor->pTree[0];  //Every time you start at the root of the tree.

  while (aCompressor->pPhase<(aLength*8) && aCompressor->pLength>0)  //Until the matrix is read or the length of the original file is reached, search in the tree.
  {
    if (nodeIsLeaf(aCompressor->pNode))  //If you get to a leaf then write down the value and return to the root.
    {
      aBufferO[lLength]=nodeGetValue(aCompressor->pNode);
      lLength++;
      aCompressor->pLength--;
      aCompressor->pNode=aCompressor->pTree[0];
    } else {  //If you are at a node take one of the edges.
      unsigned char lBit=compressorReadBit(aCompressor, aBufferI);

      if (lBit==0) aCompressor->pNode=nodeGetChildL(aCompressor->pNode);
      if (lBit==1) aCompressor->pNode=nodeGetChildR(aCompressor->pNode);
    }
  }

  aCompressor->pPhase=0;

  return lLength;
}

void compressorDestructor (compressor* aCompressor)
{
  /*
    Gets rid of the compressor.
  */
  unsigned int i;

  if (aCompressor->pTable) free(aCompressor->pTable);

  if (aCompressor->pTree)
  {
    for (i=0; i<255 && aCompressor->pTree[i]!=0; i++) free(aCompressor->pTree[i]);
    free(aCompressor->pTree);
  }

  if (aCompressor->pStack)
  {
    stackDestructor(aCompressor->pStack);
    free(aCompressor->pStack);
  }
}
